import json
from typing import List, Union
from volcenginesdkcore.signv4 import SignerV4

from ai_engine.core.llm.volcengine.ark_service import ArkService
from ai_engine.core.llm.volcengine.ark_common import EmbModel, RawData, HOST, EMBEDDING_V2_PATH, REGION_CN_NORTH1, SERVICE


class ArkEmbeddingsService(ArkService):
    """
     向量模型
    """

    def create(self,
               emb_model: EmbModel,
               raw_data: Union[RawData, List[RawData]]):

        data = self._build_data_param(raw_data)
        req_param = {
            "model": {"model_name": emb_model.model_name, "params": emb_model.params},
            "data": data
        }
        body = json.dumps(req_param)
        headers = {
            "Accept": "application/json",
            "Content-Type": "application/json",
            "Host": HOST
        }
        # 生成签名
        SignerV4.sign(path=EMBEDDING_V2_PATH,
                      method="POST",
                      headers=headers,
                      body=body,
                      query={},
                      ak=self.ak,
                      sk=self.sk,
                      region=REGION_CN_NORTH1,
                      service=SERVICE)

        resp = self.json_exception(url='https://{}{}'.format(HOST, EMBEDDING_V2_PATH),
                                   headers=headers,
                                   body=body)
        res = json.loads(resp)
        return res["data"]

    def _build_data_param(self, raw_data: Union[RawData, List[RawData]]):
        data = []
        if isinstance(raw_data, RawData):
            param = {"data_type": raw_data.data_type}
            if raw_data.text != "":
                param["text"] = raw_data.text
            if raw_data.image != "":
                param["image"] = raw_data.image
            data.append(param)
        elif isinstance(raw_data, List):
            for item in raw_data:
                param = {"data_type": item.data_type}
                if item.text != "":
                    param["text"] = item.text
                if item.image != "":
                    param["image"] = item.image
                data.append(param)
        return data


if __name__ == "__main__":
    ak = "AKLTYTZmZjU0NjEwYmYzNDhlNGI2NTZiN2E4YjBlODhjM2M"
    sk = "TURrME1EZGlPVGszWlRjeU5EUmlNV0l4WVRZeE9UVTBaREUxTURVM1pqYw=="
    #
    # list = [RawData("text", "hello1"), RawData("text", "hello2")]
    # embeddingsv2 = ArkEmbeddingsApi(ak=ak, sk=sk)
    #
    # res2 = embeddingsv2.create(EmbModel("bge-m3", params={"return_token_usage": True}), list)
    # print("res2=" + str(res2))

    # list = [RawData("text", "hello1"), RawData("text", "hello2")]
    # vikingdb_service = VikingDBService()
    # res = vikingdb_service.embedding_v2(EmbModel("bge-m3", params={"return_token_usage": True}), list)
    # print(res)
